#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
使用训练好的BERT模型预测账单分类
"""

import torch
import os
from transformers import BertTokenizer

# 设置环境变量使用HF镜像源
os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'

def load_model(model_path='./model_cache/bert_transaction_classifier', use_onnx=True):
    """加载训练好的模型和tokenizer"""
    # 检查ONNX模型是否存在
    onnx_model_path = model_path + '.onnx' if not model_path.endswith('.onnx') else model_path
    if use_onnx and os.path.exists(onnx_model_path):
        import onnxruntime as ort
        tokenizer = BertTokenizer.from_pretrained(model_path if not model_path.endswith('.onnx') else './model_cache/bert_transaction_classifier')
        model = ort.InferenceSession(onnx_model_path, providers=['CPUExecutionProvider'])
        return model, tokenizer, 'onnx'
    
    # 如果ONNX模型不存在，则使用PyTorch模型
    if not os.path.exists(model_path):
        raise FileNotFoundError(f"模型路径 {model_path} 不存在，请先训练模型或检查路径。")
    
    from transformers import BertForSequenceClassification
    tokenizer = BertTokenizer.from_pretrained(model_path)
    # 尝试加载模型，优先使用safetensors格式
    try:
        model = BertForSequenceClassification.from_pretrained(model_path, use_safetensors=True)
    except Exception as e:
        # 如果safetensors格式不可用，则使用默认格式
        model = BertForSequenceClassification.from_pretrained(model_path)
    return model, tokenizer, 'pytorch'


def predict_category(text, model, tokenizer, model_type='pytorch', max_length=128):
    """预测文本的分类"""
    # 定义一级分类标签
    labels = ['餐饮', '交通', '购物', '娱乐', '医疗', '教育', '住房', '其他']
    
    # 编码输入文本
    encoding = tokenizer.encode_plus(
        text,
        add_special_tokens=True,
        max_length=max_length,
        padding='max_length',
        truncation=True,
        return_tensors='pt'
    )
    
    # 预测
    if model_type == 'onnx':
        # 使用ONNX模型进行推理
        inputs = {
            'input_ids': encoding['input_ids'].cpu().numpy(),
            'attention_mask': encoding['attention_mask'].cpu().numpy()
        }
        outputs = model.run(None, inputs)
        logits = torch.tensor(outputs[0])
        predictions = torch.nn.functional.softmax(logits, dim=-1)
        predicted_class_id = predictions.argmax().item()
        confidence = predictions[0][predicted_class_id].item()
    else:
        # 使用PyTorch模型进行推理
        model.eval()
        with torch.no_grad():
            outputs = model(**encoding)
            predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
            predicted_class_id = predictions.argmax().item()
            confidence = predictions[0][predicted_class_id].item()
    
    return labels[predicted_class_id], confidence


def predict_category_with_subcategory(text, model, tokenizer, model_type='pytorch', max_length=128):
    """预测文本的一级和二级分类
    
    注意：这是一个简化的实现，实际应用中可能需要更复杂的逻辑
    来处理二级分类预测
    """
    # 先预测一级分类
    primary_category, confidence = predict_category(text, model, tokenizer, model_type, max_length)
    
    # 这里应该有更复杂的逻辑来预测二级分类
    # 为了演示，我们使用简单的规则
    subcategory_rules = {
        '餐饮': ['快餐', '正餐', '饮品', '零食'],
        '交通': ['公共交通', '打车', '自驾', '火车飞机'],
        '购物': ['服饰', '电子产品', '日用品', '家居'],
        '娱乐': ['电影', '游戏', '旅游', '运动'],
        '医疗': ['药品', '检查', '治疗', '保健品'],
        '教育': ['书籍', '培训', '学费', '文具'],
        '住房': ['房租', '水电费', '物业费', '装修'],
        '其他': ['礼品', '捐赠', '意外支出', '其他']
    }
    
    # 简单的二级分类预测逻辑（实际应用中应该使用另一个模型）
    import random
    subcategories = subcategory_rules.get(primary_category, ['其他'])
    secondary_category = random.choice(subcategories)
    
    return primary_category, secondary_category, confidence


def main():
    """主函数，用于测试分类功能"""
    try:
        # 加载模型
        model, tokenizer, model_type = load_model()
        print(f"模型加载成功！模型类型: {model_type}")
        
        # 测试分类
        test_texts = [
            "在肯德基吃午餐",
            "乘坐地铁去公司",
            "在淘宝购买衣服",
            "观看电影",
            "去医院看病"
        ]
        
        for text in test_texts:
            category, confidence = predict_category(text, model, tokenizer, model_type)
            print(f"文本: '{text}' -> 分类: {category} (置信度: {confidence:.4f})")
            
    except Exception as e:
        print(f"错误: {e}")


if __name__ == "__main__":
    main()